﻿#include <stdio.h>

#define NUM_GRID_LUT_SRC (33)    /* 17,33,65,129 */
#define NUM_GRID_LUT_DST (129)
#define MAX_GRID_LUT ((NUM_GRID_LUT_SRC>NUM_GRID_LUT_DST) ? NUM_GRID_LUT_SRC:NUM_GRID_LUT_DST)
#define IMAGE_BIT_LENGTH (10)

#define CLIP(x,l,h) (((x) < (l)) ? (l) : (((x) > (h)) ? (h) : (x)))

struct sLUTConfig {
	float* lut;
	int gridSize;
	float skip;
	int bitLength;
	int maxValue;
	int uBounds;
};

/* Cube file defines 3D-LUT and the size is normally NxNxNx3, (N=17,33,65,129) (3 is corresponding R,G,B)*/
static float CUBE[] = {
	#include "S_Night_Time_2.cube"
};

void convertRGB2YUV(float r, float g, float b, float uvofs, float yuv[3])
{
	yuv[0] = 0.299f * r + 0.587f * g + 0.114f * b;
	yuv[1] = -0.169f * r - 0.331f * g + 0.500f * b + uvofs;
	yuv[2] = 0.500f * r - 0.419f * g - 0.081f * b + uvofs;
}

void convertYUV2RGB(float yuv[], float rgb[], float uvofs, float clipL, float clipH)
{
	yuv[1] -= uvofs;
	yuv[2] -= uvofs;
	rgb[0] = CLIP(1.000f * yuv[0] + 1.402f * yuv[2], clipL, clipH);
	rgb[1] = CLIP(1.000f * yuv[0] - 0.344f * yuv[1] - 0.714f * yuv[2], clipL, clipH);
	rgb[2] = CLIP(1.000f * yuv[0] + 1.772f * yuv[1], clipL, clipH);
}

#define INTERPOLATION(pt0,pt1,wt0) ((pt0)*(wt0)+(pt1)*(1.0f-(wt0)))
#define AT(a,b,c,d,size) (3*(a*size*size+b*size+c)+d)
#define GET_RGB(lut,pt,c,size) (lut[AT(pt[2],pt[1],pt[0],c,size)])
#define GET_YUV(lut,pt,c,size) (lut[AT(pt[0], pt[1], pt[2], c, size)])
#define GET(lut,pt,c,size, cube) (cube ? GET_RGB(lut,pt,c,size) : GET_YUV(lut,pt,c,size))

void interporationCubic(float* pt, sLUTConfig* lutConfig, float finalVal[3], bool cubeConvert = false)
{
	int gridBase[3];
	float wt[3];
	for (int i = 0; i < 3; i++) {
		gridBase[i] = (int)(pt[i] / lutConfig->skip);
		wt[i] = gridBase[i] + 1.f - (pt[i] / lutConfig->skip);  /* wt=1: the point is same as base.*/
	}

	int gridCubic[8][3] = {
		{gridBase[0]    , gridBase[1]    , gridBase[2]    },
		{gridBase[0]    , gridBase[1]    , gridBase[2] + 1},
		{gridBase[0]    , gridBase[1] + 1, gridBase[2]    },
		{gridBase[0]    , gridBase[1] + 1, gridBase[2] + 1},
		{gridBase[0] + 1, gridBase[1]    , gridBase[2]    },
		{gridBase[0] + 1, gridBase[1]    , gridBase[2] + 1},
		{gridBase[0] + 1, gridBase[1] + 1, gridBase[2]    },
		{gridBase[0] + 1, gridBase[1] + 1, gridBase[2] + 1},
	};

	float val4[4][3];
	float val2[2][3];
	float* lut = lutConfig->lut;
	int size = lutConfig->gridSize;
	for (int i = 0; i < 3; i++) {
		for (int j = 0; j < 4; j++) {
			val4[j][i] = INTERPOLATION(GET(lut, gridCubic[j], i, size, cubeConvert), GET(lut, gridCubic[j + 4], i, size, cubeConvert), wt[0]);
		}
		for (int j = 0; j < 2; j++) {
			val2[j][i] = INTERPOLATION(val4[j][i], val4[j + 2][i], wt[1]);
		}
		finalVal[i] = INTERPOLATION(val2[0][i], val2[1][i], wt[2]);
	}
}

void convertLUT_RGB2YUV(sLUTConfig *LUTs, sLUTConfig *LUTconv)
{
	const int uvofset = LUTconv->uBounds >> 1;

	for (int y = 0; y < LUTconv->gridSize; y++) {
		for (int u = 0; u < LUTconv->gridSize; u++) {
			for (int v = 0; v < LUTconv->gridSize; v++) {
				float yuv[] = {y*LUTconv->skip, u* LUTconv->skip, v* LUTconv->skip };
				float rgb[3];
				float convRGB[3];
				float convYUV[3];
				convertYUV2RGB(yuv, rgb, uvofset, 0, LUTconv->maxValue);
				interporationCubic(rgb, LUTs, convRGB, true);
				convertRGB2YUV(convRGB[0]* LUTconv->maxValue, convRGB[1]* LUTconv->maxValue, convRGB[2]* LUTconv->maxValue, uvofset, convYUV);
				for (int i = 0; i < 3; i++) {
					LUTconv->lut[AT(y,u,v,i,LUTconv->gridSize)] = convYUV[i];
				}
			}
		}
	}
}

void convertPixWithLUT(int* yuv, sLUTConfig* lut, int dst[3])
{
	float fyuv[] = { (float)yuv[0], (float)yuv[1], (float)yuv[2] };
	float conv[3];
	interporationCubic(fyuv, lut, conv);

	for (int i = 0; i<3; i++){
		dst[i] = CLIP((int)(conv[i] + 0.5f), 0, lut->maxValue);
	}
}



/* Load YUV data:: Actually, we hope the YUV data to use is output from IPE. */
int*  loadImage(int* width, int* height)
{
	/***** Test image data -> ******/
	/*  */
	const int imageWidth = 2;
	const int imageHeight = 3;
	/* YUV Image  ##Because this is a test data, it has only 2x2 pixels.*/
	const int yuvTestData[imageHeight*imageWidth][3] = {
		{117,552,479}, {297,500, 530},
		{405,371,788}, {529,596, 398},
		{878,101,443}, {911,267, 103}
	};
	/***** <- Test image data ******/

	int index = 0;
	int* yuvSrc = new int[imageWidth * imageHeight * 3];
	for (int y = 0; y < imageHeight; y++) {
		for (int x = 0; x < imageWidth; x++) {
			for (int i = 0; i < 3; i++) yuvSrc[index++] = yuvTestData[y * imageWidth + x][i];
		}
	}
	*width = imageWidth;
	*height = imageHeight;
	return yuvSrc;
}


void configLUT(sLUTConfig* lut, int gridNum)
{
	lut->bitLength = IMAGE_BIT_LENGTH;
	lut->gridSize = gridNum;
	lut->uBounds = 1 << lut->bitLength;
	lut->maxValue = lut->uBounds - 1;
	lut->skip = lut->uBounds / (float)(lut->gridSize - 1);

}


void initialize(sLUTConfig* LUTs, sLUTConfig* LUTconv)
{
	configLUT(LUTs, NUM_GRID_LUT_SRC);
	LUTs->lut = CUBE;  /* Actually, cube file is read by app side and send the data to driver. */
	configLUT(LUTconv, NUM_GRID_LUT_DST);
	LUTconv->lut = new float[NUM_GRID_LUT_DST * NUM_GRID_LUT_DST * NUM_GRID_LUT_DST * 3];

	convertLUT_RGB2YUV(LUTs, LUTconv);
}

void process(int srcYUV[], int dstYUV[], int pixCount, sLUTConfig* lut)
{
	for (int i = 0; i < pixCount; i++) {
		convertPixWithLUT(&srcYUV[i*3], lut, &dstYUV[i*3]);
	}
}

int main()
{
	sLUTConfig LUTs;
	sLUTConfig LUTconv;

	/* Call once during configStream. */
	initialize(&LUTs, &LUTconv);



	/* Call every frame for every YUVs. */
	int w,h;
	int* srcYUV = loadImage(&w, &h);  /* When you implement driver, the srcYUV should be picked from IPEOUT. */
	int pixCount = w*h;
	int* dstYUV = new int[pixCount*3];
	process(srcYUV, dstYUV, w*h, &LUTconv);
	for (int i = 0; i < w * h; i++) {
		printf("{%d,%d,%d}\n", dstYUV[i * 3], dstYUV[i * 3 + 1], dstYUV[i * 3 + 2]);
	}
}

